import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import os
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def plot_radar_layer(npz_file_path: str, layer_index: int = 0, save_path: str = None):
    """
    读取雷达数据NPZ文件并绘制指定层的图像

    参数:
        npz_file_path: NPZ文件路径
        layer_index: 要绘制的层索引，默认为0（第一层）
        save_path: 图像保存路径，为None则直接显示
    """
    # 检查文件是否存在
    if not os.path.exists(npz_file_path):
        raise FileNotFoundError(f"文件不存在: {npz_file_path}")

    # 加载NPZ文件
    try:
        with np.load(npz_file_path) as data:
            # 提取所需数据
            radar_data = data['data']
            lon_grid = data['lon_grid']
            lat_grid = data['lat_grid']
            layer_heights = data['layer_heights']
            grid_size = data['grid_size_km']

            radar_data[radar_data < -10] = np.nan

            # 检查层索引是否有效
            if layer_index < 0 or layer_index >= len(radar_data):
                raise IndexError(f"层索引无效，有效范围: 0-{len(radar_data) - 1}")

            # 获取指定层数据
            layer_data = radar_data[layer_index]
            layer_height = layer_heights[layer_index]

            # 创建自定义颜色映射（适合雷达数据）
            colors = ['#0000ff', '#00ffff', '#00ff00', '#ffff00', '#ff0000', '#8b00ff']
            cmap = LinearSegmentedColormap.from_list('radar_cmap', colors, N=100)

            # 创建图形
            plt.figure(figsize=(10, 8))

            # 绘制热图
            im = plt.pcolormesh(
                lon_grid, lat_grid, layer_data,
                cmap=cmap, shading='gouraud'
            )

            # 添加颜色条
            cbar = plt.colorbar(im)
            cbar.set_label('雷达反射率/强度')

            # 设置标题和坐标轴标签
            plt.title(
                f'雷达数据 - 第{layer_index + 1}层 (高度: {layer_height:.2f}km) - 分辨率: {grid_size}km×{grid_size}km')
            plt.xlabel('经度 (°E)')
            plt.ylabel('纬度 (°N)')

            # 设置经纬度范围
            plt.xlim(np.min(lon_grid), np.max(lon_grid))
            plt.ylim(np.min(lat_grid), np.max(lat_grid))

            # 添加网格线
            plt.grid(True, linestyle='--', alpha=0.7)

            # 调整布局
            plt.tight_layout()

            # 保存或显示图像
            if save_path:
                # 创建保存目录（如果不存在）
                os.makedirs(os.path.dirname(save_path), exist_ok=True)
                plt.savefig(save_path, dpi=300, bbox_inches='tight')
                print(f"图像已保存至: {save_path}")
            else:
                plt.show()

    except Exception as e:
        print(f"处理数据时出错: {str(e)}")
        raise


# 使用示例
if __name__ == "__main__":
    # NPZ文件路径（替换为实际文件路径）
    npz_file = "processed_radar_data.npz"

    try:
        # 绘制第一层数据
        plot_radar_layer(
            npz_file_path=npz_file,
            layer_index=0,  # 第一层
            # 可选：指定保存路径
            # save_path="radar_layer_1.png"
        )

        # 如需绘制其他层，可修改layer_index参数
        # 例如绘制第三层：
        # plot_radar_layer(npz_file, layer_index=2)

    except Exception as e:
        print(f"执行出错: {str(e)}")